from interpolate import interpolate_to_1km, extract_region
from typing import Dict, Tuple, List
from read_data import parse_radar_data
import numpy as np

def select_every_other_layer(interpolated_data: Dict) -> Dict:
    """每隔一层选取数据，保留一半的层数"""
    # 确保总层数为偶数，方便每隔一层选取
    total_layers = interpolated_data['N']
    if total_layers % 2 != 0:
        print(f"警告：总层数{total_layers}不是偶数，将选取前{total_layers // 2 + 1}层中的间隔层")

    # 每隔一层选取（0, 2, 4, ...）
    selected_indices = list(range(0, total_layers, 2))

    # 构建结果
    selected_data = {
        'H': interpolated_data['H'],
        'L': interpolated_data['L'],
        'N': len(selected_indices),
        'grid_size_km': interpolated_data['grid_size_km'],
        'layer_heights': [interpolated_data['layer_heights'][i] for i in selected_indices],
        'proj_mode': interpolated_data['proj_mode'],
        'data': [interpolated_data['data'][i] for i in selected_indices],
        'lon_grid': interpolated_data['lon_grid'],
        'lat_grid': interpolated_data['lat_grid'],
        'corners': interpolated_data['corners']
    }

    return selected_data


def process_radar_data_for_ai(zip_file_path: str) -> Dict:
    """
    处理雷达数据以满足AI建模需求：
    1. 提取38-42N，114-118E范围内的数据
    2. 插值到1km×1km分辨率
    3. 从30层中每隔一层选取15层数据
    """
    # 1. 解析原始数据
    print("正在解析雷达数据...")
    parsed_data = parse_radar_data(zip_file_path)

    # 2. 提取指定区域数据 (38-42N，114-118E)
    print("正在提取指定区域数据...")
    region_data = extract_region(
        parsed_data,
        min_lat=38.0, max_lat=42.0,
        min_lon=114.0, max_lon=118.0
    )

    # 3. 插值到1km×1km分辨率
    print("正在插值到1km×1km分辨率...")
    interpolated_data = interpolate_to_1km(region_data)

    # 4. 每隔一层选取数据
    print("正在选取间隔层数据...")
    processed_data = select_every_other_layer(interpolated_data)

    print(f"数据处理完成：{processed_data['L']}×{processed_data['H']}格点，{processed_data['N']}层")
    return processed_data


# 使用示例
if __name__ == "__main__":
    # 替换为实际的雷达数据ZIP文件路径
    radar_zip_path = "D:\\Mosaic20250812134200_Z.dat.zip"

    try:
        # 处理数据
        ai_ready_data = process_radar_data_for_ai(radar_zip_path)

        # 可以在这里添加保存处理后数据的代码
        # 例如使用np.savez保存为numpy格式
        # np.savez('processed_radar_data.npz', **ai_ready_data)

        print("数据已准备好用于AI建模训练")
        print(f"数据形状：{len(ai_ready_data['data'])}层，{ai_ready_data['L']}×{ai_ready_data['H']}格点")

    except Exception as e:
        print(f"处理数据时出错：{str(e)}")